use std::collections::{BinaryHeap, HashMap};
use std::cmp::Reverse;

struct Solution;


impl Solution {
    pub fn diagonal_sort(mat: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
        let mut res = vec![];
        let row = mat.len();
        let col = mat[0].len();


        let mut map = HashMap::<i32, BinaryHeap<Reverse<i32>>>::new();

        for (i, vec) in mat.into_iter().enumerate() {
            for (j, val) in vec.into_iter().enumerate() {
                let diff = i as i32 - j as i32;

                let heap = map.get_mut(&diff);

                let heap =
                    if let Some(heap) = heap {
                        heap
                    } else {
                        let heap = BinaryHeap::new();
                        map.insert(diff, heap);
                        map.get_mut(&diff).unwrap()
                    };

                heap.push(Reverse(val));
            }
        }

        for i in 0..row {
            let mut tmp = vec![];
            for j in 0..col {
                let diff = i as i32 - j as i32;
                let val = map.get_mut(&diff).unwrap().pop().unwrap().0;
                tmp.push(val);
            }
            res.push(tmp);
        }

        res
    }
}


pub fn sort_the_matrix_diagonally(mat: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
    Solution::diagonal_sort(mat)
}

#[test]
fn test() {
    let vec = vec![vec![3, 3, 1, 1], vec![2, 2, 1, 2], vec![1, 1, 1, 2]];
    let res = sort_the_matrix_diagonally(vec);
    println!("{:?}", res);
}


fn main() {

}
